from __future__ import annotations
from typing import List, Callable, Optional, TypeVar, Generic, Set, Dict
from collections import deque
from heapq import heappop, heappush, heapreplace, heapify
from time import time

T = TypeVar('T')


class Node(Generic[T]):
    def __init__(self, state: T, parent: Optional[Node] = None, cost: float = 0.0, heuristic: float = 0.0):
        self.state: T = state
        self.parent: Optional[T] = parent
        self.cost: float = cost
        self.heuristic: float = heuristic

    def __lt__(self, other: Node) -> bool:
        return (self.cost + self.heuristic) < (other.cost + other.heuristic)


def dsf(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]]):
    # frontier is where we've yet to go
    start_t = time()
    frontier: deque = deque()
    frontier.append(Node(initial, None))
    explored: Set[T] = {initial}

    while frontier:
        current_node: Node[T] = frontier.pop()
        current_state: T = current_node.state
        if goal_test(current_state):
            print("dsf spend: %s s" % (time() - start_t))
            return current_node

        for child in successors(current_state):
            if child in explored:
                continue
            explored.add(child)
            child.parent = current_node
            frontier.append(child)
    return None


def bsf(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]]):
    start_t = time()
    frontier: deque = deque()
    frontier.append(Node(initial, None))
    explored: Set[T] = {initial}

    while frontier:
        current_node: Node[T] = frontier.popleft()
        current_state: T = current_node.state
        if goal_test(current_state):
            print("bsf spend: %s s" % (time() - start_t))
            return current_node

        for child in successors(current_state):
            if child in explored:
                continue
            explored.add(child)
            child.parent = current_node
            frontier.append(child)
    return None


def astar(initial: T, goal_test: Callable[[T], bool], successors: Callable[[T], List[T]],
          heuristic: Callable[[T], float]) -> Optional[Node[T]]:
    """
    A*算法
    :param initial:初始状态
    :param goal_test:判断是否达标
    :param successors:
    :param heuristic:
    :return:
    """
    frontier: List[Node[T]] = []
    heapify(frontier)
    heappush(frontier, Node(initial, None, 0.0, heuristic(initial)))

    explored: Dict = {initial: 0.0}
    while frontier:
        current_node: Node = heappop(frontier)
        current_state: T = current_node.state
        if goal_test(current_state):
            return current_node
        for child in successors(current_state):
            new_cost = current_node.cost + 1
            if child not in explored or explored[child] > new_cost:
                explored[child] = new_cost
                child.parent = current_node
                heappush(frontier, child)
    return None


def node_to_path(node: Node[T]) -> List[T]:
    path = [node.state]
    while node.parent is not None:
        node: Node[T] = node.parent
        path.append(node.state)
    path.reverse()
    return path
